`include "defines.v"

module MEM_Stage (
    // Jump/Branch Signals
    input  wire                         inst_jump,
    input  wire [ 2: 0]                 inst_branch_type,
    input  wire [ 3: 0]                 psw_flags,
    input  wire [`BUSLEN-1:0]           branch_target_pc_i,

    output wire                         inst_jump_or_branch_taken_flag,
    output wire [`BUSLEN-1:0]           branch_target_pc_o,

    // Regfile/CSR R/W
    input  wire [`REGFILE_ADDR_LEN-1:0] rd_w_addr_i,
    input  wire                         rd_w_en_i,
    input  wire                         rd_w_src_i,

    output wire [`REGFILE_ADDR_LEN-1:0] rd_w_addr_o,
    output wire                         rd_w_en_o,
    output wire                         rd_w_src_o,

    input  wire                         csr_w_en_i,
    input  wire [`CSR_ADDR_LEN-1:0]     csr_w_addr_i,

    output wire                         csr_w_en_o,
    output wire [`CSR_ADDR_LEN-1:0]     csr_w_addr_o,

    // Memory Access
    input  wire                         mem_r_en,
    input  wire [ 2: 0]                 mem_r_ext_type_i,
    input  wire                         mem_w_en,
    input  wire [ 1: 0]                 mem_w_size_i,
    input  wire [`BUSLEN-1:0]           ALU_out_i,
    input  wire [`BUSLEN-1:0]           regbus_B_i,

    output wire [ 2: 0]                 mem_r_ext_type_o,
    output wire [`BUSLEN-1:0]           regbus_C,
    output wire [`BUSLEN-1:0]           regbus_D,

    output wire                         mem_stall_req,

    // To axi_rw
    // Read
	output wire                         mem_r_req,
    input  wire                         mem_r_handshaked,
	input  wire                         mem_r_okay,
    input  wire [`AXI_DATA_WIDTH-1:0]   mem_data_read,
    output wire [`AXI_ADDR_WIDTH-1:0]   mem_r_addr,
    output wire [1:0]                   mem_r_size,
    input  wire [1:0]                   mem_r_resp,

    // Write
    output wire                         mem_w_req,
    input  wire                         mem_w_okay,
    output wire [`AXI_DATA_WIDTH-1:0]   mem_data_write,
    output wire [`AXI_ADDR_WIDTH-1:0]   mem_w_addr,
    output wire [1:0]                   mem_w_size,
    input  wire [1:0]                   mem_w_resp,

    // Privileged Instruction / Exception
    input  wire                         is_system_inst,

    input  wire                         inst_ecall_i,
    input  wire                         inst_mret_i,

    output wire                         inst_ecall_o,
    output wire                         inst_mret_o,

    //sim
    output wire                         mem_okay              //TODO:delete before soc
);

    //sim
    assign mem_okay = mem_r_okay | mem_w_okay;   //TODO:delete before soc

    assign regbus_C = is_system_inst ?  ALU_out_i   :  mem_data_read  ;
    assign regbus_D = is_system_inst ?  regbus_B_i  :  ALU_out_i  ;

    //----------Memory/Regfile/CSR Access----------//
    // Memory
    assign mem_stall_req    = (mem_r_req & ~mem_r_okay & (mem_r_resp == 2'b0) | mem_w_req & ~mem_w_okay & (mem_w_resp == 2'b0));
    assign mem_r_ext_type_o = mem_r_ext_type_i;

    // Regfile
    assign rd_w_en_o  = rd_w_en_i;
    assign rd_w_src_o = rd_w_src_i;
    assign rd_w_addr_o     = rd_w_addr_i;

    // CSR
    assign csr_w_en_o   = csr_w_en_i;
    assign csr_w_addr_o = csr_w_addr_i;

    //----------To axi_rw----------//
    // Read
    assign mem_r_req  = mem_r_en & ~mem_r_okay;
    assign mem_r_addr = ALU_out_i;
    assign mem_r_size = mem_r_ext_type_i[1:0];

    // Write
    assign mem_w_req      = mem_w_en & ~mem_w_okay;
    assign mem_w_addr     = ALU_out_i;
    assign mem_data_write = regbus_B_i;
    assign mem_w_size     = mem_w_size_i;

    //----------Branch/Jump Signals----------//
    reg branch_taken;

    always @(*) begin
        case (inst_branch_type)
            `NOTBRANCH : branch_taken = 1'b0;
            `EQ        : branch_taken = `ZF;
            `NE        : branch_taken = ~`ZF;
            `LT        : branch_taken = `SF^`OF;
            `GE        : branch_taken = `SF^~`OF | `ZF;
            `LTU       : branch_taken = `CF;
            `GEU       : branch_taken = ~`CF | `ZF;
            default    : branch_taken = 1'b0;
        endcase
    end

    assign inst_jump_or_branch_taken_flag = inst_jump | branch_taken;
    assign branch_target_pc_o             = branch_target_pc_i;

    //----------Privileged Instructions----------//
    assign inst_ecall_o = inst_ecall_i;
    assign inst_mret_o  = inst_mret_i;

endmodule